/*
  Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include "analytic_case.h"
#include "blackbody.h"
#include "blackbody_grain.h"
#include "misc.h"
#include "wd01_Brent_PAH_grain_model.h"

class uniform_factory {
public:
  typedef T_grid::T_grid_impl::T_data T_data;
  typedef refinement_accuracy_data<T_data> T_racc;
  typedef cell_tracker<T_data> T_cell_tracker;

  const T_densities rho_, n_tol_;
  const T_float rmax_;
  const int n_lambda, maxlevel;
  const int n_threads_;

  T_densities get_rho(const T_cell_tracker& c) {
    // return 0 if outside of radius rmax
    T_densities rho = (dot(c.getcenter(), c.getcenter())< rmax_*rmax_) ? rho_ : T_densities(rho_*0.);
    return rho;
  };

  uniform_factory (const T_densities& rho, T_densities n_tol, 
		   int nl, int ml, T_float rmax, int nt):
    rho_ (rho), n_lambda (nl), maxlevel(ml), n_tol_(n_tol), 
    rmax_(rmax), n_threads_(nt) {};
  
  bool refine_cell_p (const T_cell_tracker& c) {
    const int level = c.code().level();
    if ((level<maxlevel)&&
	(abs(rmax_*rmax_-dot(c.getcenter(),c.getcenter())) < 
	 3*dot(c.getsize(),c.getsize()))) // && any(get_rho(c)>0))
      return true;
      
    // cell is refined to keep column density below n_tol
    T_densities rho = get_rho(c);
    T_densities n2 (dot(c.getsize(), c.getsize())*rho*rho);
    bool refine = all( n2 > (n_tol_*n_tol_));
    //if(!refine) cout << "No refinement level " << level << " at " << c.getcenter() << c.getsize() << endl;
    return refine;
  };
  bool unify_cell_p (const T_cell_tracker& c, const T_racc& racc) {
    const vec3d sz = c.getsize();
    const T_densities mean (racc.sum.get_absorber().densities()/racc.n);
    const T_densities stdev2 (racc.sumsq.get_absorber().densities()/racc.n-
			      mean*mean);

    // unify only if the resulting cell would be < n_tol
    T_densities n2 (dot(sz,sz)*mean*mean);
    bool unify = all( n2 < (n_tol_*n_tol_));
    
    return all (stdev2==0 ) && unify;
  };
  T_data get_data (const T_cell_tracker& c) {
    return T_data(T_data::T_emitter(), 
		  T_data::T_absorber(get_rho(c), vec3d(0,0,0), T_lambda(n_lambda)));
  }
    
  int n_threads () {return n_threads_;};
};

void setup (po::variables_map& opt)
{
  T_unit_map units;
  units ["length"] = "au";
  units ["mass"] = "Msun";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  const int maxlevel= opt["maxlevel"].as<int>();
  const T_float source_temp = opt["source_temp"].as<T_float>();
  const T_float source_lum = opt["source_lum"].as<T_float>();
  const int npix = opt["npix"].as<int>();
  const int n_threads = opt["n_threads"].as<int>();
  const T_float tau = opt["tau"].as<T_float>();
  const T_float tau_tol = opt["tau_tol"].as<T_float>();
  const T_float grid_extent=5;
  const T_float cameradist=1e3;
  const string grain_model = opt["grain_model"].as<string>();

  // we need a preferences object to pass parameters to grain model
  Preferences p;
  p.setValue("grain_data_directory", 
	     word_expand(opt["grain_data_dir"].as<string>())[0]);
  p.setValue("grain_model", grain_model);
  p.setValue("wd01_parameter_set", opt["wd01_parameter_set"].as<string>());
  p.setValue("template_pah_fraction", 0.5);
  p.setValue("use_dl07_opacities", true);
  p.setValue("n_threads", n_threads);
  p.setValue("use_grain_temp_lookup", true);
  p.setValue("use_cuda", false);
 
  array_1 lambda;
  lambda.reference(logspace(opt["lambda_min"].as<T_float>(),
			    opt["lambda_max"].as<T_float>(),
			    opt["n_lambda"].as<int>()));

  // load grain model
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(load_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p,units));
  T_dust_model model (sv);
  model.set_wavelength(lambda);

  std::vector<T_float> lambdas(lambda.begin(), lambda.end());
  const int lambda550 = lower_bound(lambdas.begin(), lambdas.end(),
				    0.551e-6) - lambdas.begin()-1;

  // find opacity normalization for whatever dust we're using
  const T_float unit_opacity_rho = 1.0/sv[0]->opacity()(lambda550);
  T_densities rho(1); rho=unit_opacity_rho/grid_extent*tau;
  T_densities n_tol(1); n_tol = tau_tol*unit_opacity_rho;

  // create factory
  uniform_factory factory (rho,n_tol, lambda.size(), maxlevel, 
			   grid_extent, n_threads);

  blackbody bb(source_temp, 3.91e26*source_lum);
  boost::shared_ptr<T_emission> em
    (new pointsource_emission<polychromatic_policy, T_rng_policy>
     (vec3d (.000, .00, -.00), bb.emission(lambda)));

  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (.0, 0.0) );
  //cam_pos.push_back(std:: make_pair (1.5, .8) );
  boost::shared_ptr<T_emergence> cameras
    (new T_emergence(cam_pos, cameradist, 2.1*grid_extent, npix));

  run_case(opt, 
	   factory,
	   units,
	   lambda,
	   vec3d(-grid_extent,-grid_extent,-grid_extent),
	   vec3d(grid_extent,grid_extent,grid_extent),
	   *em,
	   *cameras,
	   model); 
}

void pre_shoot(T_xfer&) {};

std::string description()
{
  return "Optically thick cloud with central point source test problem.";
}

void add_custom_options(po::options_description& desc)
{
  desc.add_options()
    ("source_temp", po::value<T_float>()->default_value(5800),
     "temperature of central blackbody source")
    ("source_lum", po::value<T_float>()->default_value(1),
     "luminosity of central source (in Lsun)")
    ("radius", po::value<T_float>()->default_value(5),
     "radius of cloud (in AU)")
    ("tau", po::value<T_float>()->default_value(100),
     "V-band optical depth to center of cloud")
    ("grain_model", po::value<string>()->default_value("wd01_Brent_PAH"), "dust grain model (use \"blackbody\" for blackbody grains)")
    ;
}

void print_custom_options(const po::variables_map& opt)
{
  std::cout
       << "source_temp = " << opt["source_temp"].as<T_float>() << '\n'
       << "source_lum = " << opt["source_lum"].as<T_float>() << '\n'
       << "radius = " << opt["radius"].as<T_float>() << '\n'
       << "tau = " << opt["tau"].as<T_float>() << '\n'
       << "grain_model = " << opt["grain_model"].as<std::string>() << '\n'
    ;
}
